"""
PointPipline require result['points'].shape = [NPoint, NDim],
"""

from mmcv.utils.misc import is_tuple_of
from numpy.lib.arraysetops import isin
from ..builder import PIPELINES
import numpy as np
import copy

def farthest_point_sample(point, npoint):
    """
    Input:
        xyz: pointcloud data, [N, D]
        npoint: number of samples
    Return:
        indexes: sampled pointcloud index, [npoint]
    """
    N, D = point.shape
    xyz = point[:,:3]
    centroids = np.zeros((npoint,))
    distance = np.ones((N,)) * 1e10
    farthest = np.random.randint(0, N)
    for i in range(npoint):
        centroids[i] = farthest
        centroid = xyz[farthest, :]
        dist = np.sum((xyz - centroid) ** 2, -1)
        distance = np.minimum(distance, dist)
        farthest = np.argmax(distance, -1)
    return centroids.astype(np.int32)

@PIPELINES.register_module()
class PointSample:
    def __init__(self, num_point, uniform=False, replace=False):
        self.num_point = num_point
        self.uniform = uniform
        self.replace = replace
        assert not (uniform == False and replace == True), "farthest_point_sample do not support replace=True"
    
    def __call__(self, results):
        """Call function to flip bounding boxes, masks, semantic segmentation
        maps.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: sampled data
        """
        points = results['points']
        if self.uniform:
            index = np.random.choice(np.arange(len(points)), self.num_point, replace=self.replace)
        else:
            index = farthest_point_sample(points, self.num_point)
        results['points'] = points[index]

        if 'part_labels' in results:
            results['part_labels'] = results['part_labels'][index]
        return results

    def __repr__(self):
        return self.__class__.__name__ + f'(num_point={self.num_point})'

@PIPELINES.register_module()
class PointNormalize:
    def __call__(self, results):
        points = results['points']
        centroid = np.mean(points, axis=0)
        points = points - centroid
        m = np.max(np.sqrt(np.sum(points**2, axis=1)))
        points = points / m
        results['points'] = points
        return results

@PIPELINES.register_module()
class PointDropout:
    def __init__(self, max_ratio=0.875):
        self.max_ratio = max_ratio

    def __call__(self, results):
        points = results['points']
        ratio =  np.random.random() * self.max_ratio # 0~0.875
        drop_idx = np.where(np.random.random((points.shape[0]))<=ratio)[0]
        if len(drop_idx)>0:
            points[drop_idx, :] = points[0, :] # set to the first point
        results['points'] = points

        if 'seg_labels' in results:
            seg_labels = results['seg_labels']
            if len(drop_idx) > 0:
                seg_labels[drop_idx, :] = seg_labels[0, :]
            results['seg_labels'] = seg_labels
        
        return results


@PIPELINES.register_module()
class PointScale:
    def __init__(self, range=(0.67, 1.5)):
        assert is_tuple_of(range, float) and len(range) == 2
        self.range = range
    
    def __call__(self, results):
        d = results['points'].shape[-1]
        scale = np.random.uniform(self.range[0], self.range[1], size=[d])
        results['points'] *= scale
        return results


@PIPELINES.register_module()
class PointShift:
    def __init__(self, range=0.2):
        assert isinstance(range, float)
        self.range = range
    
    def __call__(self, results):
        d = results['points'].shape[-1]
        shifts = np.random.uniform(-self.range, self.range, size=[d])
        results['points'] += shifts
        return results


@PIPELINES.register_module()
class PointNoise:
    def __init__(self, range=0.05):
        assert isinstance(range, float)
        self.range = range
    
    def __call__(self, results):
        points = results['points']
        N, C = points.shape
        pmin, pmax = points.min(-1), points.max(-1)
        noise_range = (pmax - pmin) * self.range
        noises = []
        for i in range(C):
            noise = np.random.uniform(-noise_range[i]/2, noise_range[i]/2, N).astype(np.float32)
            noises.append(noise)
        noises = np.stack(noises, axis=-1)
        points = points + noises
        results['points'] = points.astype(np.float32)
        return results

@PIPELINES.register_module()
class PointShuffle:
    def __init__(self):
        pass

    def __call__(self, results):
        results = copy.deepcopy(results)
        np.random.shuffle(results['points'])
        return results